package gov.ornl.stucco.structured;

import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InterruptedIOException;
import java.nio.channels.FileLock;
import java.nio.channels.OverlappingFileLockException;
import java.io.File;
import java.util.HashMap;
import java.util.List;
import java.util.Map;


import gov.ornl.stucco.ConfigLoader;
import gov.ornl.stucco.RabbitMQConsumer;
import gov.ornl.stucco.GraphConstructor;
import gov.ornl.stucco.AlignFactory;
import gov.ornl.stucco.Align;
import gov.ornl.stucco.preprocessors.PreprocessSTIX;
import gov.ornl.stucco.preprocessors.PreprocessSTIX.Vertex;
import gov.ornl.stucco.stix_extractors.BugtraqExtractor;
import gov.ornl.stucco.stix_extractors.CaidaExtractor;
import gov.ornl.stucco.stix_extractors.CIF1d4Extractor;
import gov.ornl.stucco.stix_extractors.CIFZeusTrackerExtractor;
import gov.ornl.stucco.stix_extractors.CIFEmergingThreatsExtractor;
import gov.ornl.stucco.stix_extractors.CleanMxVirusExtractor;
import gov.ornl.stucco.stix_extractors.ClientBannerExtractor;
import gov.ornl.stucco.stix_extractors.CpeExtractor;
import gov.ornl.stucco.stix_extractors.CveExtractor;
import gov.ornl.stucco.stix_extractors.FSecureExtractor;
import gov.ornl.stucco.stix_extractors.GeoIPExtractor;
import gov.ornl.stucco.stix_extractors.HoneExtractor;
import gov.ornl.stucco.stix_extractors.LoginEventExtractor;
import gov.ornl.stucco.stix_extractors.MalwareDomainListExtractor;
import gov.ornl.stucco.stix_extractors.MetasploitExtractor;
import gov.ornl.stucco.stix_extractors.NvdToStixExtractor;
import gov.ornl.stucco.stix_extractors.PackageListExtractor;
import gov.ornl.stucco.stix_extractors.ServiceListExtractor;
import gov.ornl.stucco.stix_extractors.SophosExtractor;
import gov.ornl.stucco.graph_extractors.ArgusGraphExtractor;
import gov.ornl.stucco.graph_extractors.HTTPDataGraphExtractor;
import gov.ornl.stucco.graph_extractors.HTTPRDataGraphExtractor;
import gov.ornl.stucco.graph_extractors.SituGraphExtractor;
import gov.ornl.stucco.graph_extractors.SnoGraphExtractor;
import gov.ornl.stucco.graph_extractors.DNSRecordGraphExtractor;
import gov.ornl.stucco.graph_extractors.BannerGraphExtractor;
import gov.pnnl.stucco.doc_service_client.DocServiceClient;
import gov.pnnl.stucco.doc_service_client.DocServiceException;
import gov.pnnl.stucco.doc_service_client.DocumentObject;

import org.mitre.stix.stix_1.STIXPackage;
import org.json.JSONObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.rabbitmq.client.GetResponse;

public class StructuredTransformer {
    private static final Logger logger = LoggerFactory.getLogger(StructuredTransformer.class);
    private static final String[] argusHeaders = {"StartTime", "Flgs", "Proto", "SrcAddr", "Sport", "Dir", "DstAddr", "Dport", "TotPkts", "TotBytes", "State"};
    private RabbitMQConsumer consumer;
    private DocServiceClient docClient;
    private PreprocessSTIX preprocessSTIX;
    private GraphConstructor constructGraph;
    private Align alignment;
    private boolean persistent;
    private boolean outputToDB;
    private boolean outputToSTIXFile;
    private String outputSTIXPath;
    private final String HOSTNAME_KEY = "hostName";

    public StructuredTransformer() {
        logger.info("loading config file from default location");
        ConfigLoader configLoader = new ConfigLoader();
        init(configLoader);
    }

    private void init(ConfigLoader configLoader) {
        Map<String, Object> configMap;
        String exchange = null;
        String queue = null;
        String host = null;
        int port = -1;
        String user = null;
        String password = null;
        String[] bindingKeys = null;
        outputToDB = true;
        outputToSTIXFile = false;
        outputSTIXPath = "stixoutput.xml";
        try {
            configMap = configLoader.getConfig("structured_data");
            outputToDB = Boolean.parseBoolean(String.valueOf(configMap.get("outputToDB")));
            outputToSTIXFile = Boolean.parseBoolean(String.valueOf(configMap.get("outputToSTIXFile")));
            outputSTIXPath = String.valueOf(configMap.get("outputSTIXPath"));
            exchange = String.valueOf(configMap.get("exchange"));
            queue = String.valueOf(configMap.get("queue"));
            host = String.valueOf(configMap.get("host"));
            port = Integer.parseInt(String.valueOf(configMap.get("port")));
            user = String.valueOf(configMap.get("username"));
            password = String.valueOf(configMap.get("password"));
            persistent = Boolean.parseBoolean(String.valueOf(configMap.get("persistent")));
            @SuppressWarnings("unchecked")
            List<String> bindings = (List<String>) (configMap.get("bindings"));
            bindingKeys = new String[bindings.size()];
            bindingKeys = bindings.toArray(bindingKeys);
        } catch (FileNotFoundException e1) {
            logger.error("Error loading configuration.", e1);
            System.exit(-1);
        } catch (Exception e) {
            logger.error("Error parsing configuration.", e);
            System.exit(-1);
        }
        logger.info("Config file loaded and parsed");
        try {
            logger.info("Connecting to rabbitMQ with this info: \nhost: " + host + "\nport: " + port +
                    "\nexchange: " + exchange + "\nqueue: " + queue +
                    "\nuser: " + user + "\npass: " + password);
            consumer = new RabbitMQConsumer(exchange, queue, host, port, user, password, bindingKeys);
            consumer.openQueue();
        } catch (IOException e) {
            logger.error("Error initializing RabbitMQ connection.", e);
            System.exit(-1);
        }
        logger.info("RabbitMQ connected.");
        try {
            preprocessSTIX = new PreprocessSTIX();
            constructGraph = new GraphConstructor();
            alignment = AlignFactory.getAlign();
            logger.info("DB connection created.  Connecting to document service...");
            configMap = configLoader.getConfig("document_service");
            host = String.valueOf(configMap.get("host"));
            port = Integer.parseInt(String.valueOf(configMap.get("port")));
            docClient = new DocServiceClient(host, port);
        } catch (IOException e) {
            logger.error("Error initializing Alignment and/or DB connection.", e);
            System.exit(-1);
        }
        logger.info("Alignment obj, DB connection, and Document service client created.  Initialization complete!");
    }

    public void run() {
        GetResponse response = null;
        boolean fatalError = false;
        do {
            try {
                response = consumer.getMessage();
            } catch (IOException e) {
                logger.error("Encountered RabbitMQ IO error:", e);
                fatalError = true;
            }
            while (response != null && !fatalError) {
                String routingKey = response.getEnvelope().getRoutingKey().toLowerCase();
                long deliveryTag = response.getEnvelope().getDeliveryTag();
                String message;
                if (response.getBody() != null) {
                    message = new String(response.getBody());
                    boolean contentIncluded = false;
                    Map<String, Object> headerMap = response.getProps().getHeaders();
                    if ((headerMap != null) && (headerMap.containsKey("HasContent")))
                        contentIncluded = Boolean.valueOf(String.valueOf(headerMap.get("HasContent")));
                    System.out.println("Recieved: " + routingKey + " deliveryTag=[" + deliveryTag + "] message- " + message);
                    String content = message;
                    if (!contentIncluded && !routingKey.endsWith(".sophos") && !routingKey.endsWith(".bugtraq")) {
                        String docId = content.trim();
                        try {
                            DocumentObject document = docClient.fetch(docId);
                            String rawContent = document.getDataAsString();
                            JSONObject jsonContent = new JSONObject(rawContent);
                            content = (String) jsonContent.get("document");
                        } catch (DocServiceException e) {
                            logger.error("Could not fetch document '" + docId + "' from Document-Service.", e);
                        } catch (Exception e) {
                            logger.error("Other error in handling document '" + docId + "' from Document-Service.", e);
                        }
                    }
                    String docIDs = null;
                    if (!contentIncluded)
                        docIDs = message;
                    Map<String, String> metaDataMap = null;
                    if (routingKey.endsWith(".hone")) {
                        if ((headerMap != null) && (headerMap.containsKey(HOSTNAME_KEY))) {
                            metaDataMap = new HashMap<>();
                            String hostname = String.valueOf(headerMap.get(HOSTNAME_KEY));
                            metaDataMap.put(HOSTNAME_KEY, hostname);
                        }
                    }
                    if (outputToDB) {
                        JSONObject graph = generateGraph(routingKey, content, metaDataMap, docIDs);
                        if (graph != null)
                            alignment.load(graph);
                    }
                    if (outputToSTIXFile) {
                        System.out.println("通过解析结构化数据来构建STIX内容");
                        STIXPackage stixPackage = generateSTIX(routingKey, content, metaDataMap, docIDs);
                        System.out.println("输出解析结构化数据内容到文件");
                        String stixContent = null;
                        if (stixPackage == null) {
                            JSONObject graph = generateGraph(routingKey, content, metaDataMap, docIDs);
                            JSONObject vertices = graph.optJSONObject("vertices");
                            if (vertices != null) {
                                StringBuilder str = new StringBuilder();
                                for (String id : vertices.keySet()) {
                                    str.append(vertices.getJSONObject(id).getString("sourceDocument"));
                                    str.append("\n");
                                }
                                stixContent = str.toString();
                            }
                        } else {
                            stixContent = stixPackage.toXMLString(true);
                        }
                        try {
                            FileOutputStream fos = new FileOutputStream(new File(outputSTIXPath), true);
                            try {
                                boolean written = false;
                                do {
                                    try {
                                        // Lock it!
                                        FileLock lock = fos.getChannel().lock();
                                        try {
                                            // Write the bytes.
                                            fos.write(stixContent.getBytes());
                                            written = true;
                                        } finally {
                                            // Release the lock.
                                            lock.release();
                                        }
                                    } catch (Exception ofle) {
                                        System.out.println(ofle.getMessage());
                                    }
                                } while (!written);
                            } catch (Exception ex) {
                                logger.warn("Failed to lock " + outputSTIXPath, ex);
                            }
                            fos.close();
                        } catch (IOException e) {
                            logger.error("Could not write stix xml file: ", e);
                            fatalError = true;
                        }
                    }
                    try {
                        consumer.messageProcessed(deliveryTag);
                    } catch (IOException e) {
                        logger.error("Encountered RabbitMQ IO error:", e);
                        fatalError = true;
                    }
                } else
                    try {
                        consumer.retryMessage(deliveryTag);
                    } catch (IOException e) {
                        logger.error("Encountered RabbitMQ IO error:", e);
                        fatalError = true;
                    }
                try {
                    response = consumer.getMessage();
                } catch (IOException e) {
                    fatalError = true;
                }
            }
        } while (persistent && !fatalError);
        try {
            consumer.close();
        } catch (IOException e) {
            logger.error("Encountered RabbitMQ IO error when closing connection:", e);
        }
    }


    /**
     * @param routingKey  determines which extractor to use
     * @param content     the text to parse
     * @param metaDataMap any additional required info, which is not included in the content
     * @param docIDs      if the content is from the document server, this is its id(s).  Only included for debugging output.
     * @return
     */
    private JSONObject generateGraph(String routingKey, String content, Map<String, String> metaDataMap, String docIDs) {
        boolean stixDocument = false;
        STIXPackage stixPackage = null;
        JSONObject graph = null;
        try {
            if (routingKey.endsWith(".argus")) {
                ArgusGraphExtractor extractor = new ArgusGraphExtractor(argusHeaders, content);
                return extractor.getGraph();
            } else if (routingKey.endsWith(".http")) {
                //TODO: find name of http file ... for now (for testing) it just has .http extencion
                HTTPDataGraphExtractor httpExtractor = new HTTPDataGraphExtractor(content);
                return httpExtractor.getGraph();
            } else if (routingKey.endsWith(".httpr")) {
                //TODO: find name of http file ... for now (for testing) it just has .http extencion
                HTTPRDataGraphExtractor httprExtractor = new HTTPRDataGraphExtractor(content);
                return httprExtractor.getGraph();
            } else if (routingKey.endsWith(".situ")) {
                SituGraphExtractor situExtractor = new SituGraphExtractor(content);
                return situExtractor.getGraph();
            } else if (routingKey.endsWith(".sno")) {
                SnoGraphExtractor snoExtractor = new SnoGraphExtractor(content);
                return snoExtractor.getGraph();
            } else if (routingKey.endsWith(".dnsrecord")) {
                DNSRecordGraphExtractor dnsExt = new DNSRecordGraphExtractor(content);
                return dnsExt.getGraph();
            } else if (routingKey.endsWith(".serverbanner")) {
                BannerGraphExtractor bannerExt = new BannerGraphExtractor(content);
                return bannerExt.getGraph();
            } else if (routingKey.endsWith(".stix")) {
                stixDocument = true;
            } else {
                stixPackage = generateSTIX(routingKey, content, metaDataMap, docIDs);
            }
            if (stixPackage != null) {
                Map<String, Vertex> stixElements = preprocessSTIX.normalizeSTIX(stixPackage.toXMLString());
                graph = constructGraph.constructGraph(stixElements);
            } else if (stixDocument) {
                Map<String, Vertex> stixElements = preprocessSTIX.normalizeSTIX(content);
                graph = constructGraph.constructGraph(stixElements);
            } else {
                logger.warn("Unexpected null stix package for routing key '" + routingKey + "'.");
            }
        } catch (RuntimeException e) {
            logger.error("Error occurred with routingKey = " + routingKey);
            logger.error("										docIDs = " + docIDs);
            logger.error("										content = " + content);
            return null;
        }
        return graph;
    }

    /**
     * @param routingKey  determines which extractor to use
     * @param content     the text to parse
     * @param metaDataMap any additional required info, which is not included in the content
     * @param docIDs      if the content is from the document server, this is its id(s).  Only included for debugging output.
     * @return
     */
    private STIXPackage generateSTIX(String routingKey, String content, Map<String, String> metaDataMap, String docIDs) {
        boolean stixDocument = false;
        STIXPackage stixPackage = null;
        JSONObject graph = null;
        try {
            if (routingKey.endsWith(".cve")) {
                CveExtractor cveExtractor = new CveExtractor(content);
                stixPackage = cveExtractor.getStixPackage();
            } else if (routingKey.endsWith(".nvd")) {
                NvdToStixExtractor nvdExt = new NvdToStixExtractor(content);
                stixPackage = nvdExt.getStixPackage();
            } else if (routingKey.endsWith(".cpe")) {
                CpeExtractor cpeExtractor = new CpeExtractor(content);
                stixPackage = cpeExtractor.getStixPackage();
            } else if (routingKey.endsWith(".maxmind")) {
                GeoIPExtractor geoIPExtractor = new GeoIPExtractor(content);
                stixPackage = geoIPExtractor.getStixPackage();
            } else if (routingKey.endsWith(".metasploit")) {
                MetasploitExtractor metasploitExtractor = new MetasploitExtractor(content);
                stixPackage = metasploitExtractor.getStixPackage();
            } else if (routingKey.replaceAll("\\-", "").endsWith(".cleanmx")) {
                CleanMxVirusExtractor virusExtractor = new CleanMxVirusExtractor(content);
                stixPackage = virusExtractor.getStixPackage();
            } else if (routingKey.endsWith(".login_events")) {
                LoginEventExtractor loginEventExtractor = new LoginEventExtractor(content);
                stixPackage = loginEventExtractor.getStixPackage();
            } else if (routingKey.endsWith(".installed_package")) {
                PackageListExtractor packageListExtractor = new PackageListExtractor(content);
                stixPackage = packageListExtractor.getStixPackage();
            } else if (routingKey.endsWith(".1d4")) {
                CIF1d4Extractor cifExtractor = new CIF1d4Extractor(content);
                stixPackage = cifExtractor.getStixPackage();
            } else if (routingKey.endsWith(".zeustracker")) {
                CIFZeusTrackerExtractor cifExtractor = new CIFZeusTrackerExtractor(content);
                stixPackage = cifExtractor.getStixPackage();
            } else if (routingKey.endsWith(".emergingthreats")) {
                CIFEmergingThreatsExtractor cifExtractor = new CIFEmergingThreatsExtractor(content);
                stixPackage = cifExtractor.getStixPackage();
            } else if (routingKey.endsWith(".servicelist")) {
                ServiceListExtractor serviceListExtractor = new ServiceListExtractor(content);
                stixPackage = serviceListExtractor.getStixPackage();
            } else if (routingKey.endsWith(".clientbanner")) {
                ClientBannerExtractor clientBannerExtractor = new ClientBannerExtractor(content);
                stixPackage = clientBannerExtractor.getStixPackage();
            } else if (routingKey.replaceAll("\\-", "").endsWith(".fsecure")) {
                FSecureExtractor fSecureExt = new FSecureExtractor(content);
                stixPackage = fSecureExt.getStixPackage();
            } else if (routingKey.endsWith(".malwaredomainlist")) {
                MalwareDomainListExtractor mdlExt = new MalwareDomainListExtractor(content);
                stixPackage = mdlExt.getStixPackage();
            } else if (routingKey.endsWith(".hone")) {
                HoneExtractor honeExtractor = null;
                if ((metaDataMap != null) && (metaDataMap.containsKey(HOSTNAME_KEY))) {
                    honeExtractor = new HoneExtractor(content, metaDataMap.get(HOSTNAME_KEY));
                } else {
                    honeExtractor = new HoneExtractor(content);
                }
                stixPackage = honeExtractor.getStixPackage();
            } else if (routingKey.endsWith(".caida")) {
                //TODO: ensure file names match
                String as2org = null;
                String pfx2as = null;
                String[] items = content.split("\\r?\\n");
                for (String item : items) {
                    String docId = item.split("\\s+")[0];
                    String sourceURL = item.split("\\s+")[1];
                    String rawItemContent = null;
                    String itemContent = null;
                    try {
                        DocumentObject document = docClient.fetch(docId);
                        rawItemContent = document.getDataAsString();
                        JSONObject jsonContent = new JSONObject(rawItemContent);
                        itemContent = (String) jsonContent.get("document");
                    } catch (DocServiceException e) {
                        logger.error("Could not fetch document '" + docId + "' from Document-Service. URL was: " + sourceURL, e);
                        logger.error("Complete message content was:\n" + content);
                        return null;
                    }
                    if (sourceURL.contains("as2org")) {
                        as2org = itemContent;
                    } else if (sourceURL.contains("pfx2as")) {
                        pfx2as = itemContent;
                    } else {
                        logger.warn("unexpected URL (sophos) " + sourceURL);
                    }
                }
                if (as2org != null && pfx2as != null) {
                    CaidaExtractor caidaExtractor = new CaidaExtractor(as2org, pfx2as);
                    stixPackage = caidaExtractor.getStixPackage();
                }
            } else if (routingKey.endsWith(".sophos")) {
                String summary = null;
                String details = null;
                String[] items = content.split("\\r?\\n");
                for (String item : items) {
                    String docId = item.split("\\s+")[0];
                    String sourceURL = item.split("\\s+")[1];
                    String rawItemContent = null;
                    String itemContent = null;
                    try {
                        DocumentObject document = docClient.fetch(docId);
                        rawItemContent = document.getDataAsString();
                        JSONObject jsonContent = new JSONObject(rawItemContent);
                        itemContent = (String) jsonContent.get("document");
                    } catch (DocServiceException e) {
                        logger.error("Could not fetch document '" + docId + "' from Document-Service. URL was: " + sourceURL, e);
                        logger.error("Complete message content was:\n" + content);
                        return null;
                    }
                    if (sourceURL.contains("/detailed-analysis.aspx")) {
                        details = itemContent;
                    } else if (sourceURL.contains(".aspx")) {
                        summary = itemContent;
                    } else {
                        logger.warn("unexpected URL (sophos) " + sourceURL);
                    }
                }
                if (summary != null && details != null) {
                    SophosExtractor sophosExt = new SophosExtractor(summary, details);
                    stixPackage = sophosExt.getStixPackage();
                } else {
                    logger.warn("Sophos: some required fields were null, skipping group.\nMessage was:" + content);
                }
            } else if (routingKey.endsWith(".bugtraq")) {
                String info = null;
                String discussion = null;
                String exploit = null;
                String solution = null;
                String references = null;
                String[] items = content.split("\\r?\\n");
                for (String item : items) {
                    String docId = item.split("\\s+")[0];
                    String sourceURL = item.split("\\s+")[1];
                    String rawItemContent = null;
                    String itemContent = null;
                    try {
                        DocumentObject document = docClient.fetch(docId);
                        rawItemContent = document.getDataAsString();
                        JSONObject jsonContent = new JSONObject(rawItemContent);
                        itemContent = (String) jsonContent.get("document");
                    } catch (DocServiceException e) {
                        logger.error("Could not fetch document '" + docId + "' from Document-Service. URL was: " + sourceURL, e);
                        logger.error("Complete message content was:\n" + content);
                        return null;
                    }
                    if (sourceURL.contains("/info")) {
                        info = itemContent;
                    } else if (sourceURL.contains("/discuss")) { //interestingly, "/discuss" and "/discussion" are both valid urls for this item
                        discussion = itemContent;
                    } else if (sourceURL.contains("/exploit")) {
                        exploit = itemContent;
                    } else if (sourceURL.contains("/solution")) {
                        solution = itemContent;
                    } else if (sourceURL.contains("/references")) {
                        references = itemContent;
                    } else {
                        logger.warn("unexpected URL (bugtraq) " + sourceURL);
                    }
                }
                if (info != null && discussion != null && exploit != null && solution != null && references != null) {
                    BugtraqExtractor bugtraqExt = new BugtraqExtractor(info, discussion, exploit, solution, references);
                    stixPackage = bugtraqExt.getStixPackage();
                } else {
                    logger.warn("Bugtraq: some required fields were null, skipping group.\nMessage was:" + content);
                    if (docIDs != null) {
                        logger.error("Problem docid(s):\n" + docIDs);
                    }
                }
            } else {
                logger.warn("Unexpected routing key encountered '" + routingKey + "'.\n"
                        + "\t(If running with 'outputToSTIXFile: true', this source "
                        + "may be using a 'graph extractor' instead of a 'stix extractor'/)");
            }
        } catch (RuntimeException e) {
            logger.error("Error occurred with routingKey = " + routingKey);
            logger.error("										docIDs = " + docIDs);
            logger.error("										content = " + content);
            e.printStackTrace();
            return null;
        }

        return stixPackage;
    }

    /**
     * @param args
     */
    public static void main(String[] args) {
        System.setProperty("STUCCO_DB_TYPE", "POSTGRESQL");
        System.setProperty("STUCCO_DB_CONFIG", "postgresql.yml");
        System.setProperty("STUCCO_DB_INDEX_CONFIG", "");
        StructuredTransformer structProcess;
        structProcess = new StructuredTransformer();
        structProcess.run();
    }
}
